#!/usr/bin/env python
# -*- coding:utf-8 -*-
# __author__ = 'mahaojing'

import commands
import sys
import json

#################### render starts
# {% if cni.type == "vpccni" %}

network_id = "{{ cni.vpccni.network_id }}"
tenant_id = "{{ cni.vpccni.customerTenantId }}"
cluster_nodes = [
{% for i in BM_HOSTS %}
    {
        "subnet_id": "{{ hostvars[i].vpccni_subnet_id }}",
        "node_hostname": "{{ hostvars[i].inventory_hostname }}"
    },
{% endfor %}
]
# {% else %}

network_id = ""
tenant_id = ""
cluster_nodes = {}
# {% endif %}

cni_type = "{{ cni.type }}"
#################### render ends
# {% raw %}
bootstrap_token_ttl = '24h'
vpcctl_conf_file = '/etc/kubernetes/vpcctl.conf'
vpc_init_ports_conf = '/etc/kubernetes/init-ports.conf'


class K8SNode(object):
    def __init__(self, node_name, subnet_id):
        self.node_name = node_name
        self.subnet_id = subnet_id


class K8SInit(object):
    def __init__(self, cluster, kubeconfig, nodes):
        self.cluster = cluster
        self.kubeconfig = kubeconfig
        self.nodes = nodes
        self.diff_nodes = []

    def create_bootstrap_token(self):
        cmd = 'kubeadm token create --kubeconfig %s --ttl %s' % (self.kubeconfig, bootstrap_token_ttl)
        status, output = commands.getstatusoutput(cmd)
        if status != 0:
            self._failed_exit('Failed to create bootstrap token: %s' % output)
        print('Generated Bootstrap Token(TTL:%s): %s' % (bootstrap_token_ttl, output))

    def init_vpcport(self):
        self.generate_vpcctl_conf()
        self.vpcctl_init()
        self.save_vpcports_to_k8s()

    def force_init_vpcport(self):
        self.diff_nodes = self.nodes
        self._write_vpcctl_conf()
        self.vpcctl_init()
        self.save_vpcports_to_k8s()

    def generate_vpcctl_conf(self):
        self._get_diff_nodes()
        self._write_vpcctl_conf()
        print('generate /etc/kubernetes/vpcctl.conf done.')

    def vpcctl_init(self):
        status, output = commands.getstatusoutput('/opt/cni/bin/vpcctl init')
        if status != 0:
            self._failed_exit("vpcctl init failed: %s" % output)
        print('/opt/cni/bin/vpcctl init successfully.')

    def save_vpcports_to_k8s(self):
        with open(vpc_init_ports_conf) as f:
            pydata = json.load(f)
        for vpc_port_info in pydata:
            node_hostname = vpc_port_info['binding:host_id']
            if self._node_name_in_nodes(node_hostname, self.diff_nodes):
                print('start to apply: %s' % node_hostname)
                vpc_port_info_str = json.dumps(vpc_port_info, indent=8)
                node_vpcport_str='''
apiVersion: node.netease.com/v1
kind: VpcPort
metadata:
  name: %s-%s
  namespace: kube-system
spec:
  vpcPortInfo: |
    %s''' % (self.cluster, node_hostname, vpc_port_info_str)
                node_vpcport_str = node_vpcport_str.rstrip('}') + '    }'
                cmd = "echo '%s' | kubectl apply --kubeconfig %s -f -" % (node_vpcport_str, self.kubeconfig)
                status, output = commands.getstatusoutput(cmd)
                if status != 0:
                    self._failed_exit("failed to apply %s, output: %s" % (node_hostname, output))
        print('save vpc port to k8s successfully.')

    def _failed_exit(self, message):
        print(message)
        sys.exit(1)

    def _write_vpcctl_conf(self):
        print("diff_nodes: %s" % [i.node_name for i in self.diff_nodes])
        with open(vpcctl_conf_file) as f:
            pydata = json.load(f)
        pydata['initHosts'] = []
        for n in self.diff_nodes:
            if n.node_name:
                pydata['initHosts'].append({
                    'network_id': network_id,
                    'subnet_id': n.subnet_id,
                    'node_hostname': n.node_name,
                    'tenant_id': tenant_id
                })
        write_data = json.dumps(pydata, indent=4)
        with open(vpcctl_conf_file, 'w') as f:
            f.write(write_data)

    def _get_diff_nodes(self):
        cmd = "kubectl get no --kubeconfig %s -o=go-template --template '{{range .items}}{{.metadata.name}} {{end}}'" % self.kubeconfig
        status, output = commands.getstatusoutput(cmd)
        if status != 0:
            self._failed_exit("can not get nodes: %s" % output)
        all_nodes = output.strip().split()
        diff_nodes_by_ips = {}
        self.diff_nodes = []
        for node in self.nodes:
            if node.node_name not in all_nodes:
                self.diff_nodes.append(node)
                if diff_nodes_by_ips.get(node.subnet_id, None) is None:
                    diff_nodes_by_ips.update({node.subnet_id: [node.node_name]})
                else:
                    diff_nodes_by_ips[node.subnet_id].append(node.node_name)

    def _node_name_in_nodes(self, node_name, nodes):
        for n in nodes:
            if n.node_name == node_name:
                return True
        return False


def usage():
    usage_str = """usage:
    <cluster-name>: generate vpcports and create bootstrap token of <cluster-name>
    <cluster-name> <node-hostname>: generate vpcport for node <node-hostname> of <cluster-name>
    <cluster-name> token: create bootstrap token of <cluster-name>
    """
    print(usage_str)
    sys.exit(1)


def main():
    if len(sys.argv) < 2:
        usage()
    cluster_name = sys.argv[1]
    kubeconfig = "/etc/kubernetes/%s/admin.conf" % cluster_name
    if cni_type != "vpccni" or (len(sys.argv) == 3 and sys.argv[2] == 'token'):
        k8s_init_object = K8SInit(cluster=cluster_name, kubeconfig=kubeconfig, nodes=[])
        k8s_init_object.create_bootstrap_token()
        sys.exit()
    nodes = []
    if len(sys.argv) >= 3 :
        nodes_from_argv = set(sys.argv[2:])
        for node_obj in cluster_nodes:
            if node_obj['node_hostname'] in nodes_from_argv:
                nodes.append(K8SNode(node_name=node_obj['node_hostname'], subnet_id=node_obj['subnet_id']))
        if len(nodes) != len(nodes_from_argv):
            print('can not find some nodes in cluster_nodes')
            sys.exit(1)
        k8s_init_object = K8SInit(cluster=cluster_name, kubeconfig=kubeconfig, nodes=nodes)
        k8s_init_object.force_init_vpcport()
        sys.exit()

    # base init
    for node_obj in cluster_nodes:
        nodes.append(K8SNode(node_name=node_obj['node_hostname'], subnet_id=node_obj['subnet_id']))
    k8s_init_object = K8SInit(cluster=cluster_name, kubeconfig=kubeconfig, nodes=nodes)
    k8s_init_object.init_vpcport()
    k8s_init_object.create_bootstrap_token()


if __name__ == '__main__':
    main()
# {% endraw %}